import glob
import json
import os.path

import xml.sax
from shutil import copyfile

import cv2
import numpy as np

from utils.plots import plot_one_box

try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET

seg_list = ['可行驶区域', '背景', 'Puddle', 'pit']
# "MiningTruck", "OtherEquipment", "Car", "Person", "Obstacle", "Stone", "Sign"
det_list = ['矿区运载卡车', '工程车辆', '乘用车辆', '人', '其他独立障碍物', '石头', '标志牌']
labels_paths = glob.glob("/home/rhq/xiangmu/data/train/old/*/xml/*.xml")

# images_path = "/media/hnu/Backup Plus/qipanjinVideo/0819/image1/"
images_path = "/home/rhq/xiangmu/data/train/old/*/"

palette = np.random.randint(0, 255, size=(5, 3))
palette[0] = [0, 0, 0]
palette[2] = [0, 0, 0]
palette[1] = [0, 255, 0]
palette[3] = [255, 0, 0]
palette[4] = [0, 0, 255]

colors = [[np.random.randint(0, 255) for _ in range(3)] for _ in det_list]

det_num, seg_num = 0, 0

for label_name in labels_paths:
    tree = ET.parse(label_name)
    root = tree.getroot()
    # print(root.tag)
    # print(root.attrib)
    # filename = root.findtext('filename')
    filename = label_name.split(os.sep)[-1].replace('.xml', '.jpg')  # 图片名称
    images_dir = glob.glob(os.path.join(images_path, filename))  # 图片文件所在地方
    print(filename)
    if len(images_dir) == 0:
        continue
    image = cv2.imread(images_dir[0])
    H, W, c = image.shape
    bbox, cls = [], []
    masks = []
    for object in root.iter('object'):
        classname = object.findtext('name').split(" ")[0]

        points = []

        if classname in seg_list:
            if classname == "背景":
                print(seg_list.index(classname))
            mask = np.zeros((H, W), np.uint8)
            for bndbox in object.iter('bndbox'):
            #     for i in range(1,1000):
            #      a = 'x'+str(i)
            #      b = 'y'+str(i)
                x = float(bndbox.findtext('xmin'))
                y = float(bndbox.findtext('ymin'))
                points.append([x, y])
            seg_num += 1
            cv2.fillPoly(mask, np.array([points], dtype=np.int32), seg_list.index(classname) + 1)
            masks.append(mask)
        else:
            for bndbox in object.iter('bndbox'):
                xmin = float(bndbox.findtext('xmin'))
                ymin = float(bndbox.findtext('ymin'))
                xmax = float(bndbox.findtext('xmax'))
                ymax = float(bndbox.findtext('ymax'))

            det_num += 1
            bbox.append([xmin, ymin, xmax, ymax])
            cls.append(classname)
    if len(masks) == 0:
        masks.append(np.zeros((H, W), np.int8))
    mask = np.array(masks).max(axis=0)
    mask[mask == 2] = 0
    mask[mask == 3] = 1
    mask[mask == 4] = 0
    color_seg = np.zeros((H, W, 3), dtype=np.uint8)
    for label, color in enumerate(palette):
        color_seg[mask == label, :] = color
    image[mask != 0] = image[mask != 0] * 0.5 + color_seg[mask != 0] * 0.5
    image = image.astype(np.uint8)

    save_path = "/home/rhq/xiangmu/datasets/kuangka/train"

    with open(os.path.join(save_path, "labels", filename.replace(".jpg", ".txt")), 'w') as fp:
        for name, box in zip(cls, bbox):
            plot_one_box(box, image, colors[det_list.index(name)], label=name)
            x, y, w, h = (box[0] + box[2]) / 2 / W, (box[1] + box[3]) / 2 / H, (box[2] - box[0]) / W, (
                        box[3] - box[1]) / H
            data_str = str(det_list.index(name)) + ' ' + " ".join(str(i) for i in [x, y, w, h]) + "\n"
            fp.write(data_str)

    mask_path = os.path.join(save_path, "masks", filename.replace(".jpg", ".png"))
    cv2.imwrite(mask_path, mask)
    dst_image_path = os.path.join(save_path, "images", filename)
    copyfile(images_dir[0], dst_image_path)
    cv2.imwrite(os.path.join("/home/rhq/xiangmu/datasets/kuangka/train/show/", filename), image)

print("seg num: ", seg_num, "det num: ", det_num)
